#!/usr/bin/python3
# -*- coding:UTF-8 -*-

import AutoExecUtils
import os
import re
import traceback
import argparse
from bson import ObjectId
from bson.json_util import dumps, loads


def genInsRel(db):
    # 生成应用实例调用应用实例和DB实例的关系数据，存放于REF_INS和REF_DB下
    insCollection = db["COLLECT_INS"]
    relSrcCollection = db["RELATION_INS_NETCONN"]
    relDestCollection = db["RELATION_INS_NETCONN"]
    dbInsCollection = db["COLLECT_DBINS"]
    clusterCollection = db["COLLECT_CLUSTER"]

    for ins in relSrcCollection.find(
        {"_OBJ_CATEGORY": "INS"},
        {
            "_OBJ_CATEGORY": 1,
            "_OBJ_TYPE": 1,
            "OS_ID": 1,
            "MGMT_IP": 1,
            "PORT": 1,
            "BIND": {
                "$filter": {
                    "input": "$BIND",
                    "cond": {"$not": {"$regexMatch": {"input": "$$this", "regex": "^127\.|^::1:"}}},
                }
            },
            "PEER": 1,
            "LOCAL_PEER": 1,
        },
    ).batch_size(500):
        try:
            # print(dumps(ins))
            refIns = []
            refSvc = []
            refDbs = []
            refCls = []

            mgmtIp = ins.get("MGMT_IP")
            peer = ins.get("PEER", [])
            localPeer = ins.get("LOCAL_PEER", [])
            if peer or localPeer:
                # 计算调用的应用实例
                for rel in relDestCollection.find(
                    {
                        "$or": [
                            {"BIND": {"$in": peer}},
                            {"MGMT_IP": mgmtIp, "BIND": {"$in": localPeer}},
                        ]
                    },
                    {
                        "OS_ID": 1,
                        "MGMT_IP": 1,
                        "PORT": 1,
                        "BIND": {
                            "$filter": {
                                "input": "$BIND",
                                "cond": {"$not": {"$regexMatch": {"input": "$$this", "regex": "^127\.|^::1:"}}},
                            }
                        },
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                    },
                ).batch_size(500):
                    if rel["_OBJ_CATEGORY"] == "DBINS":
                        for dbIns in dbInsCollection.find(
                            {"MGMT_IP": rel["MGMT_IP"], "PORT": rel["PORT"]},
                            {
                                "OS_ID": 1,
                                "MGMT_IP": 1,
                                "PORT": 1,
                                "INSTANCE_NAME": 1,
                                "_OBJ_CATEGORY": 1,
                                "_OBJ_TYPE": 1,
                            },
                        ).batch_size(500):
                            refInfo = {
                                "_OBJ_CATEGORY": dbIns["_OBJ_CATEGORY"],
                                "_OBJ_TYPE": dbIns["_OBJ_TYPE"],
                                "OS_ID": dbIns["OS_ID"],
                                "MGMT_IP": dbIns["MGMT_IP"],
                                "PORT": dbIns["PORT"],
                                "INSTANCE_NAME": dbIns["INSTANCE_NAME"],
                            }
                            refDbs.append(refInfo)

                            for bindPeer in rel["BIND"]:
                                (peerIp, peerPort) = bindPeer.split(":", 1)
                                refSvcInfo = {
                                    "_OBJ_CATEGORY": dbIns["_OBJ_CATEGORY"],
                                    "_OBJ_TYPE": "Service-Ports",
                                    "OS_ID": dbIns["OS_ID"],
                                    "MGMT_IP": dbIns["MGMT_IP"],
                                    "PORT": peerPort,
                                    "INSTANCE_NAME": dbIns["INSTANCE_NAME"],
                                }
                                refSvc.append(refSvcInfo)
                    else:
                        refInfo = {
                            "_OBJ_CATEGORY": rel["_OBJ_CATEGORY"],
                            "_OBJ_TYPE": rel["_OBJ_TYPE"],
                            "OS_ID": rel["OS_ID"],
                            "MGMT_IP": rel["MGMT_IP"],
                            "PORT": rel["PORT"],
                        }
                        refIns.append(refInfo)

                        for bindPeer in rel["BIND"]:
                            (peerIp, peerPort) = bindPeer.split(":", 1)
                            refSvcInfo = {
                                "_OBJ_CATEGORY": rel["_OBJ_CATEGORY"],
                                "_OBJ_TYPE": "Service-Ports",
                                "OS_ID": rel["OS_ID"],
                                "MGMT_IP": rel["MGMT_IP"],
                                "PORT": peerPort,
                                # "INSTANCE_NAME": rel["INSTANCE_NAME"],
                            }
                            refSvc.append(refSvcInfo)

                # 计算调用的集群
                for cluster in clusterCollection.find(
                    {
                        "$or": [
                            {"MEMBER_PEER": {"$in": peer}},
                            {"MGMT_IP": mgmtIp, "MEMBER_PEER": {"$in": localPeer}},
                        ]
                    },
                    {
                        "UNIQUE_NAME": 1,
                        "VIP": 1,
                        "PRIMARY_IP": 1,
                        "PORT": 1,
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                        "CLUSTER_MODE": 1,
                        "CLUSTER_SOFTWARE": 1,
                    },
                ).batch_size(500):
                    refInfo = {
                        "_OBJ_CATEGORY": cluster["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": cluster["_OBJ_TYPE"],
                        "UNIQUE_NAME": cluster["UNIQUE_NAME"],
                        "PRIMARY_IP": cluster.get("PRIMARY_IP"),
                        "VIP": cluster.get("VIP"),
                        "PORT": cluster.get("PORT"),
                        # "CLUSTER_MODE": cluster.get("CLUSTER_MODE"),
                        # "CLUSTER_SOFTWARE": cluster.get("CLUSTER_SOFTWARE"),
                    }
                    refCls.append(refInfo)
            else:
                print(
                    "WARN: Object({}/{} {}:{}) Key 'PEER' not defined.".format(
                        ins.get("_OBJ_CATEGORY"),
                        ins.get("_OBJ_TYPE"),
                        ins.get("MGMT_IP"),
                        ins.get("PORT"),
                    )
                )

            # 计算属于哪个集群
            bind = ins.get("BIND")
            if bind is not None:
                belongCls = []
                for cluster in clusterCollection.find(
                    {"MEMBER_PEER": {"$in": bind}},
                    {
                        "UNIQUE_NAME": 1,
                        "VIP": 1,
                        "PRIMARY_IP": 1,
                        "PORT": 1,
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                        "CLUSTER_MODE": 1,
                        "CLUSTER_SOFTWARE": 1,
                    },
                ).batch_size(500):
                    belongCl = {
                        "_OBJ_CATEGORY": cluster["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": cluster["_OBJ_TYPE"],
                        "UNIQUE_NAME": cluster["UNIQUE_NAME"],
                        "PRIMARY_IP": cluster.get("PRIMARY_IP"),
                        "VIP": cluster.get("VIP"),
                        "PORT": cluster.get("PORT"),
                        "CLUSTER_MODE": cluster.get("CLUSTER_MODE"),
                        "CLUSTER_SOFTWARE": cluster.get("CLUSTER_SOFTWARE"),
                    }
                    belongCls.append(belongCl)
            else:
                print(
                    "WARN: Object({}/{} {}:{}) Key 'BIND' not defined.".format(
                        ins.get("_OBJ_CATEGORY"),
                        ins.get("_OBJ_TYPE"),
                        ins.get("MGMT_IP"),
                        ins.get("PORT"),
                    )
                )

            insCollection.update_many(
                {"MGMT_IP": ins["MGMT_IP"], "PORT": ins["PORT"]},
                {
                    "$set": {
                        "BELONG_CLUSTER": belongCls,
                        "REF_DB": refDbs,
                        "REF_INS": refIns,
                        "REF_SERVICE": refSvc,
                        "REF_CLUSTER": refCls,
                    }
                },
            )

        except Exception as ex:
            print("ERROR: Unknow Error, {}".format(traceback.format_exc()))


def genDBInsRel(db):
    # 生成DB的集群关系数据，存放在DB实例的REF_CLUSTER下
    dbCollection = db["COLLECT_DBINS"]
    relSrcCollection = db["RELATION_INS_NETCONN"]
    clusterCollection = db["COLLECT_CLUSTER"]

    for db in relSrcCollection.find(
        {"_OBJ_CATEGORY": "DBINS"},
        {
            "_OBJ_CATEGORY": 1,
            "_OBJ_TYPE": 1,
            "OS_ID": 1,
            "MGMT_IP": 1,
            "PORT": 1,
            "BIND": {
                "$filter": {
                    "input": "$BIND",
                    "cond": {"$not": {"$regexMatch": {"input": "$$this", "regex": "^127\.|^::1:"}}},
                }
            },
            "PEER": 1,
            "LOCAL_PEER": 1,
        },
    ).batch_size(500):
        try:
            # print(dumps(ins))
            # 计算DB调用哪个DB集群
            refCls = []
            belongCls = []

            mgmtIp = db.get("MGMT_IP")
            peer = db.get("PEER", [])
            localPeer = db.get("LOCAL_PEER", [])
            if peer or localPeer:
                for cluster in clusterCollection.find(
                    {
                        "$or": [
                            {"MEMBER_PEER": {"$in": peer}},
                            {"MGMT_IP": mgmtIp, "MEMBER_PEER": {"$in": localPeer}},
                        ]
                    },
                    {
                        "UNIQUE_NAME": 1,
                        "VIP": 1,
                        "PRIMARY_IP": 1,
                        "PORT": 1,
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                        "CLUSTER_MODE": 1,
                        "CLUSTER_SOFTWARE": 1,
                    },
                ).batch_size(500):
                    refInfo = {
                        "_OBJ_CATEGORY": cluster["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": cluster["_OBJ_TYPE"],
                        "UNIQUE_NAME": cluster["UNIQUE_NAME"],
                        "PRIMARY_IP": cluster.get("PRIMARY_IP"),
                        "VIP": cluster.get("VIP"),
                        "PORT": cluster.get("PORT"),
                        "CLUSTER_MODE": cluster.get("CLUSTER_MODE"),
                        "CLUSTER_SOFTWARE": cluster.get("CLUSTER_SOFTWARE"),
                    }
                    refCls.append(refInfo)
            else:
                print(
                    "WARN: Object({}/{} {}:{}) Key 'PEER' not defined.".format(
                        db.get("_OBJ_CATEGORY"),
                        db.get("_OBJ_TYPE"),
                        db.get("MGMT_IP"),
                        db.get("PORT"),
                    )
                )

            # 计算DB属于哪个集群
            bind = db.get("BIND")
            if bind is not None:
                for cluster in clusterCollection.find(
                    {"MEMBER_PEER": {"$in": bind}},
                    {
                        "UNIQUE_NAME": 1,
                        "VIP": 1,
                        "PRIMARY_IP": 1,
                        "PORT": 1,
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                        "CLUSTER_MODE": 1,
                        "CLUSTER_SOFTWARE": 1,
                    },
                ).batch_size(500):
                    belongCl = {
                        "_OBJ_CATEGORY": cluster["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": cluster["_OBJ_TYPE"],
                        "UNIQUE_NAME": cluster["UNIQUE_NAME"],
                        "PRIMARY_IP": cluster.get("PRIMARY_IP"),
                        "VIP": cluster.get("VIP"),
                        "PORT": cluster.get("PORT"),
                        "CLUSTER_MODE": cluster.get("CLUSTER_MODE"),
                        "CLUSTER_SOFTWARE": cluster.get("CLUSTER_SOFTWARE"),
                    }
                    belongCls.append(belongCl)
            else:
                print(
                    "WARN: Object({}/{} {}:{}) Key 'BIND' not defined.".format(
                        db.get("_OBJ_CATEGORY"),
                        db.get("_OBJ_TYPE"),
                        db.get("MGMT_IP"),
                        db.get("PORT"),
                    )
                )

            dbCollection.update_one(
                {"MGMT_IP": db["MGMT_IP"], "PORT": db["PORT"]},
                {"$set": {"BELONG_CLUSTER": belongCls, "REF_CLUSTER": refCls}},
            )
        except Exception as ex:
            print("ERROR: Unknow Error, {}".format(traceback.format_exc()))


def genHostRel(db):
    # 生成OS和交换机的关系，存放于OS的属性ETH_INTERFACES.REMOTE_PORTS下
    hostCollection = db["COLLECT_HOST"]
    fcSwCollection = db["COLLECT_FCDEV"]
    hostUpdCollection = db["COLLECT_HOST"]
    macTblCollection = db["RELATION_NET_MACTABLE"]

    for host in hostCollection.find({}, {"MGMT_IP": 1, "HBA_INTERFACES": 1, "ETH_INTERFACES": 1}).batch_size(500):
        try:
            # 计算网卡连接的交换机端口信息
            nics = host.get("ETH_INTERFACES", [])
            refSwitchesMap = {}
            if not nics:
                print("WARN: No eth interfaces found for os:{}".format(host.get("MGMT_IP")))
                nics = []
            for nic in nics:
                # macCount = sys.maxsize
                remotePort = macTblCollection.find_one(
                    {
                        "MAC_TABLE": {
                            "$elemMatch": {
                                "MACS": nic["MAC"],
                                # 对于虚拟化的物理机接入的交换机端口会看到所有虚拟机的MAC，数量不为一
                                # "MAC_COUNT": 1,
                                # 接入层端口没有邻居信息
                                "$or": [
                                    {"NEIGHBOR_COUNT": None},
                                    {"NEIGHBOR_COUNT": 0},
                                ],
                                "PORT": {"$ne": None, "$regex": r"/"},
                            }
                        }
                    },
                    {
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                        "MGMT_IP": 1,
                        "DEV_NAME": 1,
                        "SN": 1,
                        "MAC_TABLE": {
                            "$elemMatch": {
                                "MACS": nic["MAC"],
                                "PORT": {"$ne": None, "$regex": r"/"},
                                "MAC_COUNT": 1,
                            }
                        },
                    },
                )
                refSwPorts = []
                refSwPort = {}
                if remotePort is not None and "MAC_TABLE" in remotePort:
                    refSwPort = {
                        "_OBJ_CATEGORY": remotePort["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": "SWITCH-PORT",
                        "MGMT_IP": remotePort["MGMT_IP"],
                        # "DEV_NAME": remotePort.get("DEV_NAME"),
                        "SN": remotePort.get("SN"),
                        "PORT": remotePort["MAC_TABLE"][0]["PORT"],
                    }
                    refSwPorts.append(refSwPort)

                    refSwitch = {
                        "_OBJ_CATEGORY": remotePort["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": remotePort["_OBJ_TYPE"],
                        "MGMT_IP": remotePort["MGMT_IP"],
                        "SN": remotePort["SN"],
                    }
                    refSwitchesMap[remotePort["MGMT_IP"]] = refSwitch

                nic["REMOTE_PORTS"] = refSwPorts

            # 计算HBA卡连接的光交
            refFcSwitchesMap = {}
            hbas = host.get("HBA_INTERFACES", [])
            if hbas is None:
                print("INFO: No HBA interfaces found for host:{}".format(host.get("MGMT_IP")))
                hbas = []
            for hba in hbas:
                wwpn = hba.get("WWPN")
                if wwpn is None:
                    continue

                remoteFcPort = fcSwCollection.find_one(
                    {"LINK_TABLE": {"$elemMatch": {"PEER_WWPN": wwpn, "LINK_COUNT": 1}}},
                    {
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                        "MGMT_IP": 1,
                        "DEV_NAME": 1,
                        "SN": 1,
                        "LINK_TABLE": {"$elemMatch": {"PEER_WWPN": wwpn, "LINK_COUNT": 1}},
                    },
                )
                refFcPorts = []
                refFcPort = {}
                if remoteFcPort is not None and "LINK_TABLE" in remoteFcPort:
                    refFcPort = {
                        "_OBJ_CATEGORY": remoteFcPort["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": "FCSWITCH-PORT",
                        "MGMT_IP": remoteFcPort["MGMT_IP"],
                        "SN": remoteFcPort.get("SN"),
                        "PORT": remoteFcPort["LINK_TABLE"][0]["PORT_NAME"],
                    }
                    refFcPorts.append(refFcPort)

                    refFcSwitch = {
                        "_OBJ_CATEGORY": remoteFcPort["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": remoteFcPort["_OBJ_TYPE"],
                        "MGMT_IP": remoteFcPort["MGMT_IP"],
                        "SN": remoteFcPort["SN"],
                    }
                    refFcSwitchesMap[remoteFcPort["MGMT_IP"]] = refFcSwitch

                hba["REMOTE_PORTS"] = refFcPorts

            hostUpdCollection.update_one(
                {"MGMT_IP": host["MGMT_IP"]},
                {
                    "$set": {
                        "ETH_INTERFACES": nics,
                        "HBA_INTERFACES": hbas,
                        "REF_SWITCHS": list(refSwitchesMap.values()),
                        "REF_FCSWITCHS": list(refFcSwitchesMap.values()),
                    }
                },
            )
        except Exception as ex:
            print("ERROR: Unknow Error, {}".format(traceback.format_exc()))


def genLBRel(db):
    # 生成LB和交换机的关系，存放于LB的属性PORTS.REMOTE_PORTS下
    lbCollection = db["COLLECT_LOADBALANCER"]
    lbUpdCollection = db["COLLECT_LOADBALANCER"]
    macTblCollection = db["RELATION_NET_MACTABLE"]
    relDestCollection = db["RELATION_INS_NETCONN"]

    for lb in lbCollection.find({}, {"MGMT_IP": 1, "PORTS": 1}).batch_size(500):
        try:
            # 计算网卡连接的交换机端口信息
            refSwitchesMap = {}
            ports = lb.get("PORTS", [])
            if not ports:
                print("WARN: No eth interfaces found for os:{}".format(lb.get("MGMT_IP")))
            for port in ports:
                # macCount = sys.maxsize
                remotePort = macTblCollection.find_one(
                    {
                        "MAC_TABLE": {
                            "$elemMatch": {
                                "MACS": port["MAC"],
                                # 对于虚拟化的物理机接入的交换机端口会看到所有虚拟机的MAC，数量不为一
                                # "MAC_COUNT": 1,
                                # 接入层端口没有邻居信息
                                "$or": [
                                    {"NEIGHBOR_COUNT": None},
                                    {"NEIGHBOR_COUNT": 0},
                                ],
                                "PORT": {"$ne": None, "$regex": "Aggregation"},
                            }
                        }
                    },
                    {
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                        "MGMT_IP": 1,
                        "DEV_NAME": 1,
                        "SN": 1,
                        "MAC_TABLE": {
                            "$elemMatch": {
                                "MACS": port["MAC"],
                                # 不知道这逻辑对不对，但现在可以看到结果
                                "PORT": {"$ne": None, "$regex": "Aggregation"},
                                "MAC_COUNT": {"$gt": 1},
                            }
                        },
                    },
                )
                refSwPorts = []
                refSwPort = {}
                if remotePort is not None and "MAC_TABLE" in remotePort:
                    refSwPort = {
                        "_OBJ_CATEGORY": remotePort["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": "SWITCH-PORT",
                        "MGMT_IP": remotePort["MGMT_IP"],
                        # "DEV_NAME": remotePort.get("DEV_NAME"),
                        "SN": remotePort.get("SN"),
                        "PORT": remotePort["MAC_TABLE"][0]["PORT"],
                    }
                    refSwPorts.append(refSwPort)

                    refSwitch = {
                        "_OBJ_CATEGORY": remotePort["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": remotePort["_OBJ_TYPE"],
                        "MGMT_IP": remotePort["MGMT_IP"],
                        "SN": remotePort["SN"],
                    }
                    refSwitchesMap[remotePort["MGMT_IP"]] = refSwitch

                port["REMOTE_PORTS"] = refSwPorts

            lbUpdCollection.update_one({"MGMT_IP": lb["MGMT_IP"]}, {"$set": {"PORTS": ports, "REF_SWITCHES": list(refSwitchesMap.values())}})
        except Exception as ex:
            print("ERROR: Unknow Error, {}".format(traceback.format_exc()))

    # POOL member关联instance
    for lb in lbCollection.find({}, {"MGMT_IP": 1, "VIRTUAL_SERVERS": 1}).batch_size(50):
        try:
            # 计算网卡连接的交换机端口信息
            virtualServers = lb.get("VIRTUAL_SERVERS", [])
            for vs in virtualServers:
                for pool in vs.get("POOL", []):
                    for member in pool.get("MEMBERS", []):
                        ip = member.get("IP")
                        port = member.get("PORT")
                        insInfo = relDestCollection.find_one(
                            {"$or": [{"BIND": "%s:%s" % (ip, port)}]},
                            {
                                "OS_ID": 1,
                                "MGMT_IP": 1,
                                "PORT": 1,
                                "_OBJ_CATEGORY": 1,
                                "_OBJ_TYPE": 1,
                            },
                        )
                        if insInfo is not None:
                            refInfo = {
                                "_OBJ_CATEGORY": insInfo["_OBJ_CATEGORY"],
                                "_OBJ_TYPE": insInfo["_OBJ_TYPE"],
                                "OS_ID": insInfo["OS_ID"],
                                "MGMT_IP": insInfo["MGMT_IP"],
                                "PORT": insInfo["PORT"],
                            }
                            member["REF_INS"] = [refInfo]

            lbUpdCollection.update_one({"MGMT_IP": lb["MGMT_IP"]}, {"$set": {"VIRTUAL_SERVERS": virtualServers}})
        except Exception as ex:
            print("ERROR: Unknow Error, {}".format(traceback.format_exc()))


def genOSRel(db):
    # 生成OS和交换机的关系，存放于OS的属性ETH_INTERFACES.REMOTE_PORTS下
    osCollection = db["COLLECT_OS"]
    hostCollection = db["COLLECT_HOST"]
    macTblCollection = db["RELATION_NET_MACTABLE"]
    fcSwCollection = db["COLLECT_FCDEV"]
    osUpdCollection = db["COLLECT_OS"]
    clusterCollection = db["COLLECT_CLUSTER"]

    for os in osCollection.find(
        {},
        {
            "OS_ID": 1,
            "MGMT_IP": 1,
            "PRODUCT_UUID": 1,
            "ETH_INTERFACES": 1,
            "HBA_INTERFACES": 1,
            "NFS_INFO": 1,
            "DISKS": 1,
            "HOST_ON": 1,
        },
    ).batch_size(500):
        try:
            # 计算host_on在哪个物理机上
            productId = os.get("PRODUCT_UUID")
            hostOn = os.get("HOST_ON", [])
            hostOnKvm = os.get("HOST_ON_KVM", [])
            if productId is not None:
                hostInfo = hostCollection.find_one(
                    {"GUESTOS_UUIDS": productId},
                    {"_OBJ_CATEGORY": 1, "_OBJ_TYPE": 1, "MGMT_IP": 1, "DEV_NAME": 1, "SN": 1},
                )
                if hostInfo is not None:
                    hostOn = [hostInfo]
                    hostOnKvm = [{"_OBJ_CATEGORY": "HOST", "_OBJ_TYPE": "KVM", "UUID": productId, "IP": hostInfo.get("MGMT_IP")}]

            # 计算网卡连接的交换机端口信息
            refSwitchesMap = {}
            nics = os.get("ETH_INTERFACES", [])
            if not nics:
                print("WARN: No eth interfaces found for os:{}".format(os.get("MGMT_IP")))
            for nic in nics:
                # macCount = sys.maxsize
                remotePort = macTblCollection.find_one(
                    {
                        "MAC_TABLE": {
                            "$elemMatch": {
                                "MACS": nic["MAC"],
                                # 对于虚拟化的物理机接入的交换机端口会看到所有虚拟机的MAC，数量不为一
                                # "MAC_COUNT": 1,
                                # 接入层端口没有邻居信息
                                "$or": [
                                    {"NEIGHBOR_COUNT": None},
                                    {"NEIGHBOR_COUNT": 0},
                                ],
                                "PORT": {"$ne": None, "$regex": r"/"},
                            }
                        }
                    },
                    {
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                        "MGMT_IP": 1,
                        "DEV_NAME": 1,
                        "SN": 1,
                        "MAC_TABLE": {
                            "$elemMatch": {
                                "MACS": nic["MAC"],
                                "PORT": {"$ne": None, "$regex": r"/"},
                                "MAC_COUNT": 1,
                            }
                        },
                    },
                )

                refSwPorts = []
                refSwPort = {}
                if remotePort is not None and "MAC_TABLE" in remotePort:
                    refSwPort = {
                        "_OBJ_CATEGORY": remotePort["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": "SWITCH-PORT",
                        "MGMT_IP": remotePort["MGMT_IP"],
                        # "DEV_NAME": remotePort.get("DEV_NAME"),
                        "SN": remotePort.get("SN"),
                        "PORT": remotePort["MAC_TABLE"][0]["PORT"],
                    }
                    refSwPorts.append(refSwPort)

                    refSwitch = {
                        "_OBJ_CATEGORY": remotePort["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": remotePort["_OBJ_TYPE"],
                        "MGMT_IP": remotePort["MGMT_IP"],
                        "SN": remotePort["SN"],
                    }
                    refSwitchesMap[remotePort["MGMT_IP"]] = refSwitch

                nic["REMOTE_PORTS"] = refSwPorts

            # 计算HBA卡连接的光交
            refFcSwitchesMap = {}
            hbas = os.get("HBA_INTERFACES", [])
            if hbas is None:
                print("INFO: No HBA interfaces found for os:{}".format(os.get("MGMT_IP")))
            for hba in hbas:
                wwpn = hba.get("WWPN")
                if wwpn is None:
                    continue

                remoteFcPort = fcSwCollection.find_one(
                    {"LINK_TABLE": {"$elemMatch": {"PEER_WWPN": wwpn, "LINK_COUNT": 1}}},
                    {
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                        "MGMT_IP": 1,
                        "DEV_NAME": 1,
                        "SN": 1,
                        "LINK_TABLE": {"$elemMatch": {"PEER_WWPN": wwpn, "LINK_COUNT": 1}},
                    },
                )
                refFcPorts = []
                refFcPort = {}
                if remoteFcPort is not None and "LINK_TABLE" in remoteFcPort:
                    refFcPort = {
                        "_OBJ_CATEGORY": remoteFcPort["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": "FCSWITCH-PORT",
                        "MGMT_IP": remoteFcPort["MGMT_IP"],
                        "SN": remoteFcPort.get("SN"),
                        "PORT": remoteFcPort["LINK_TABLE"][0]["PORT_NAME"],
                    }
                    refFcPorts.append(refFcPort)

                    refFcSwitch = {
                        "_OBJ_CATEGORY": remoteFcPort["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": remoteFcPort["_OBJ_TYPE"],
                        "MGMT_IP": remoteFcPort["MGMT_IP"],
                        "SN": remoteFcPort["SN"],
                    }
                    refFcSwitchesMap[remoteFcPort["MGMT_IP"]] = refFcSwitch

                hba["REMOTE_PORTS"] = refFcPorts

            # 计算NFS的引用关系(生产上存在存储nas采集用的是管理管理，但os上采到的服务ip)
            storageCollection = db["COLLECT_STORAGE"]
            nfsMounts = os.get("NFS_INFO", [])
            if nfsMounts is None:
                print("INFO: No NFS mount found for os:{}".format(os.get("MGMT_IP")))
            for nfsMount in nfsMounts:
                nfsIp = nfsMount.get("NFS_IP")
                nfsRemotePath = nfsMount.get("REMOTE_PATH")
                if nfsIp is None:
                    continue

                remoteVolumn = storageCollection.find_one(
                    {
                        # "MGMT_IP": nfsIp,
                        "IPLIST": nfsIp,
                        # 可能有点问题，华为nas接口查回来的数据只有最外层的目录，可能要写正则
                        "VOLUMES": {"$elemMatch": {"NAME": nfsRemotePath}},
                    },
                    {
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                        "MGMT_IP": 1,
                        "DEV_NAME": 1,
                        "SN": 1,
                        "VOLUMES": {"$elemMatch": {"NAME": nfsRemotePath}},
                    },
                )

                if remoteVolumn is not None:
                    nfsMount["REMOTE_VOLUME"] = [
                        {
                            "_OBJ_CATEGORY": remoteVolumn["_OBJ_CATEGORY"],
                            "_OBJ_TYPE": "STORAGE-NFS-VOLUME",
                            # "MGMT_IP": nfsIp,
                            "MGMT_IP": remoteVolumn["MGMT_IP"],
                            "SN": remoteVolumn.get("SN"),
                            "VOLUME_NAME": remoteVolumn["VOLUMES"][0]["NAME"],
                        }
                    ]
                else:
                    nfsMount["REMOTE_VOLUME"] = []

            # 计算LUN的引用关系
            storageCollection = db["COLLECT_STORAGE"]
            disks = os.get("DISKS", [])
            if disks is None:
                print("INFO: No disks mount found for os:{}".format(os.get("MGMT_IP")))
            for disk in disks:
                wwn = disk.get("WWN")
                if wwn is None:
                    continue
                remoteLun = storageCollection.find_one(
                    {
                        "LUNS": {"$elemMatch": {"WWN": wwn}},
                    },
                    {
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                        "MGMT_IP": 1,
                        "DEV_NAME": 1,
                        "SN": 1,
                        "LUNS": {"$elemMatch": {"WWN": wwn}},
                    },
                )

                if remoteLun is not None:
                    disk["REMOTE_LUN"] = [
                        {
                            "_OBJ_CATEGORY": remoteLun["LUNS"][0]["_OBJ_CATEGORY"],
                            "_OBJ_TYPE": remoteLun["LUNS"][0]["_OBJ_TYPE"],
                            "MGMT_IP": remoteLun["MGMT_IP"],
                            "SN": remoteLun.get("SN"),
                            "WWN": remoteLun["LUNS"][0]["WWN"],
                        }
                    ]
                else:
                    disk["REMOTE_LUN"] = []

            # 计算操作系统属于哪个操作系统集群
            belongCls = []
            for cluster in clusterCollection.find(
                {"MEMBER_PEER": os["MGMT_IP"]},
                {
                    "UNIQUE_NAME": 1,
                    "VIP": 1,
                    "PRIMARY_IP": 1,
                    "PORT": 1,
                    "_OBJ_CATEGORY": 1,
                    "_OBJ_TYPE": 1,
                    "CLUSTER_MODE": 1,
                    "CLUSTER_SOFTWARE": 1,
                },
            ).batch_size(500):
                belongCl = {
                    "_OBJ_CATEGORY": cluster["_OBJ_CATEGORY"],
                    "_OBJ_TYPE": cluster["_OBJ_TYPE"],
                    "UNIQUE_NAME": cluster["UNIQUE_NAME"],
                    # 'PRIMARY_IP': cluster.get('PRIMARY_IP'),
                    # 'VIP': cluster.get('VIP'),
                    # 'PORT': cluster.get('PORT'),
                    "CLUSTER_MODE": cluster.get("CLUSTER_MODE"),
                    "CLUSTER_SOFTWARE": cluster.get("CLUSTER_SOFTWARE"),
                }
                belongCls.append(belongCl)

            osUpdCollection.update_one(
                {"MGMT_IP": os["MGMT_IP"]},
                {
                    "$set": {
                        "ETH_INTERFACES": nics,
                        "HBA_INTERFACES": hbas,
                        "REF_SWITCHES": list(refSwitchesMap.values()),
                        "REF_FCSWITCHES": list(refFcSwitchesMap.values()),
                        "NFS_INFO": nfsMounts,
                        "DISKS": disks,
                        "BELONG_CLUSTER": belongCls,
                        "HOST_ON": hostOn,
                        "HOST_ON_KVM": hostOnKvm
                    }
                },
            )
        except Exception as ex:
            print("ERROR: Unknow Error, {}".format(traceback.format_exc()))


def genStorageRel(db):
    storageCollection = db["COLLECT_STORAGE"]
    storUpdCollection = db["COLLECT_STORAGE"]
    fcSwCollection = db["COLLECT_FCDEV"]
    for storage in storageCollection.find({}, {"SN": 1, "MGMT_IP": 1, "HBA_INTERFACES": 1}).batch_size(500):
        try:
            # 计算HBA卡连接的光交
            refFcSwitchesMap = {}
            hbas = storage.get("HBA_INTERFACES", [])
            storageMgmtIp = storage.get("MGMT_IP", [])
            for hba in hbas:
                wwpn = hba.get("WWPN")
                if wwpn is None:
                    continue

                remoteFcPort = fcSwCollection.find_one(
                    {"LINK_TABLE.PEER_WWPN": wwpn, "LINK_TABLE.LINK_COUNT": 1},
                    {
                        "_OBJ_CATEGORY": 1,
                        "_OBJ_TYPE": 1,
                        "MGMT_IP": 1,
                        "SN": 1,
                        "LINK_TABLE": {"$elemMatch": {"PEER_WWPN": wwpn, "LINK_COUNT": 1}},
                    },
                )
                refFcPorts = []
                refFcPort = {}
                if remoteFcPort is not None and "LINK_TABLE" in remoteFcPort:
                    refFcPort = {
                        "_OBJ_CATEGORY": remoteFcPort["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": "FCSWITCH-PORT",
                        "MGMT_IP": remoteFcPort["MGMT_IP"],
                        "SN": remoteFcPort.get("SN"),
                        "PORT": remoteFcPort["LINK_TABLE"][0]["DOMAIN_IDX"],
                    }
                    refFcPorts.append(refFcPort)
                    refFcSwitch = {
                        "_OBJ_CATEGORY": remoteFcPort["_OBJ_CATEGORY"],
                        "_OBJ_TYPE": remoteFcPort["_OBJ_TYPE"],
                        "MGMT_IP": remoteFcPort["MGMT_IP"],
                        "SN": remoteFcPort["SN"],
                    }
                    refFcSwitchesMap[remoteFcPort["MGMT_IP"]] = refFcSwitch

                hba["REMOTE_PORTS"] = refFcPorts

                storUpdCollection.update_one(
                    {"MGMT_IP": storageMgmtIp},
                    {"$set": {"HBA_INTERFACES": hbas, "REF_FCSWITCHES": list(refFcSwitchesMap.values())}},
                )
        except Exception as ex:
            print("ERROR: Unknow Error, {}".format(traceback.format_exc()))


def genSwitchRel(db):
    # 生成交换机和交换机的关系数据，存放于PORTS.NEIGHBORS下
    cdpSrcCollection = db["COLLECT_SWITCH"]
    cdpDestCollection = db["COLLECT_SWITCH"]

    for swInfo in cdpSrcCollection.find({}, {"MGMT_IP": 1, "SN": 1, "PORTS": 1}).batch_size(500):
        try:
            needUpdate = False
            refSwitchesMap = {}
            ports = swInfo.get("PORTS", [])
            for port in ports:
                if "NEIGHBORS" not in port:
                    continue

                neighbors = port.get("NEIGHBORS", [])
                for neighbor in neighbors:
                    if "DEV_NAME" in neighbor:
                        remoteCount = 0
                        for dev in cdpDestCollection.find(
                            {
                                "$or": [
                                    {"DEV_NAME": neighbor["DEV_NAME"]},
                                    {"DEV_NAME": neighbor["DEV_NAME_1"]},
                                ]
                            },
                            {"_OBJ_CATEGORY": 1, "_OBJ_TYPE": 1, "MGMT_IP": 1, "SN": 1},
                        ).limit(2):
                            remoteCount = remoteCount + 1
                        if remoteCount == 1:
                            needUpdate = True
                            # 补充neighbor的信息，原来的neighbor中只有该端口连接邻居的DEV_NAME设备名、PORT名称、NEIGHBOR_COUNT
                            neighbor["_OBJ_CATEGORY"] = dev["_OBJ_CATEGORY"]
                            neighbor["_OBJ_TYPE"] = "SWITCH-PORT"
                            neighbor["MGMT_IP"] = dev["MGMT_IP"]
                            neighbor["SN"] = dev.get("SN")

                            refSwitch = {
                                "_OBJ_CATEGORY": dev["_OBJ_CATEGORY"],
                                "_OBJ_TYPE": dev["_OBJ_TYPE"],
                                "MGMT_IP": dev["MGMT_IP"],
                                "SN": dev["SN"],
                            }
                            refSwitchesMap[dev["MGMT_IP"]] = refSwitch

            if needUpdate:
                cdpDestCollection.update_one(
                    {"MGMT_IP": swInfo["MGMT_IP"], "SN": swInfo.get("SN")},
                    {"$set": {"PORTS": ports, "REF_SWITCHES": list(refSwitchesMap.values())}},
                )
        except Exception as ex:
            print("ERROR: Unknow Error, {}".format(traceback.format_exc()))


# 处理光交中的级联端口(ZONE的WWPN_MEMBERS)
def handleFCSwitchCacadePort(db):
    fcSwCollection = db["COLLECT_FCDEV"]

    for swInfo in fcSwCollection.find(
        {
            "$and": [
                {"CASCADE_SWITCHWWNN": {"$exists": "true"}},
                {"CASCADE_SWITCHWWNN": {"$ne": ""}},
            ]
        },
        {"MGMT_IP": 1, "SN": 1, "CASCADE_SWITCHWWNN": 1, "CONFIGS": 1, "ZONES": 1},
    ).batch_size(500):
        try:
            cascade_switchwwnn = swInfo.get("CASCADE_SWITCHWWNN")
            if cascade_switchwwnn is None:
                continue

            (
                cascade_switchSn,
                doaminIdx2wwpnMap,
                doaminIdx2peerwwpnMap,
                peerwwpn2wwpnMap,
                peerwwpn2domainIdxMap,
                peerwwpn2portMap,
            ) = getCascadeSwitchAttr(fcSwCollection, cascade_switchwwnn)
            zones = swInfo.get("ZONES", [])
            zoneNeedUpdate = False
            for zone in zones:
                wwpn_members = zone.get("WWPN_MEMBERS", [])
                new_wwpn_members = []
                for member in wwpn_members:
                    # 存在本地WWPN
                    if "WWPN" in member and member["WWPN"] is not None:
                        new_wwpn_members.append(member)
                        continue

                    # 不存在WWPN,存在DOMAIN_IDX级联端口
                    if "DOMAIN_IDX" in member and member.get("DOMAIN_IDX") is not None:
                        domain_idx = member.get("DOMAIN_IDX")
                        member["SN"] = cascade_switchSn
                        member["WWPN"] = doaminIdx2wwpnMap[domain_idx]
                        member["PEER_WWPN"] = doaminIdx2peerwwpnMap[domain_idx]
                        new_wwpn_members.append(member)
                        zoneNeedUpdate = True
                    else:
                        # 不存在WWPN,不存在DOMAIN_IDX,存在PEER_WWPN 级联端口
                        if "PEER_WWPN" in member and member.get("PEER_WWPN") is not None:
                            peerwwpn = member.get("PEER_WWPN")
                            member["SN"] = cascade_switchSn
                            if peerwwpn in peerwwpn2wwpnMap:
                                wwpn = peerwwpn2wwpnMap[peerwwpn]

                                if wwpn is None:
                                    continue
                                # 找到了PEER_WWPN对应的WWPN
                                member["WWPN"] = peerwwpn2wwpnMap[peerwwpn]
                                member["DOMAIN_IDX"] = peerwwpn2domainIdxMap[peerwwpn]
                                member["PORT"] = peerwwpn2portMap[peerwwpn]
                                zoneNeedUpdate = True
                                new_wwpn_members.append(member)

                zone["WWPN_MEMBERS"] = new_wwpn_members

            if zoneNeedUpdate:
                fcSwCollection.update_one(
                    {"MGMT_IP": swInfo["MGMT_IP"], "SN": swInfo.get("SN")},
                    {"$set": {"ZONES": zones}},
                )
        except Exception as ex:
            print("ERROR: Unknow Error, {}".format(traceback.format_exc()))

        try:
            cascade_switchwwnn = swInfo.get("CASCADE_SWITCHWWNN")
            if cascade_switchwwnn is None:
                continue

            (
                cascade_switchSn,
                doaminIdx2wwpnMap,
                doaminIdx2peerwwpnMap,
                peerwwpn2wwpnMap,
                peerwwpn2domainIdxMap,
                peerwwpn2portMap,
            ) = getCascadeSwitchAttr(fcSwCollection, cascade_switchwwnn)
            configs = swInfo.get("CONFIGS", [])
            configNeedUpdate = False
            for config in configs:
                zones = config.get("ZONES")
                for zone in zones:
                    wwpn_members = zone.get("WWPN_MEMBERS", [])
                    new_wwpn_members = []
                    for member in wwpn_members:
                        # 存在本地WWPN
                        if "WWPN" in member and member["WWPN"] is not None:
                            new_wwpn_members.append(member)
                            continue
                        # 不存在WWPN,存在DOMAIN_IDX级联端口
                        if "DOMAIN_IDX" in member and member.get("DOMAIN_IDX") is not None:
                            domain_idx = member.get("DOMAIN_IDX")
                            member["SN"] = cascade_switchSn
                            member["WWPN"] = doaminIdx2wwpnMap[domain_idx]
                            member["PEER_WWPN"] = doaminIdx2peerwwpnMap[domain_idx]
                            configNeedUpdate = True
                            new_wwpn_members.append(member)
                        else:
                            # 不存在WWPN,不存在DOMAIN_IDX,存在PEER_WWPN 级联端口
                            if "PEER_WWPN" in member and member.get("PEER_WWPN") is not None:
                                peerwwpn = member.get("PEER_WWPN")
                                member["SN"] = cascade_switchSn
                                if peerwwpn in peerwwpn2wwpnMap:
                                    wwpn = peerwwpn2wwpnMap[peerwwpn]
                                    if wwpn is None:
                                        continue
                                    # 找到了PEER_WWPN对应的WWPN
                                    member["WWPN"] = peerwwpn2wwpnMap[peerwwpn]
                                    member["DOMAIN_IDX"] = peerwwpn2domainIdxMap[peerwwpn]
                                    member["PORT"] = peerwwpn2portMap[peerwwpn]
                                    configNeedUpdate = True
                                    new_wwpn_members.append(member)

                    zone["WWPN_MEMBERS"] = new_wwpn_members
            if configNeedUpdate:
                fcSwCollection.update_one(
                    {"MGMT_IP": swInfo["MGMT_IP"], "SN": swInfo.get("SN")},
                    {"$set": {"CONFIGS": configs}},
                )
        except Exception as ex:
            print("ERROR: Unknow Error, {}".format(traceback.format_exc()))


# 根据光交wwnn查询对应的光交属性
def getCascadeSwitchAttr(fcSwCollection, switchwwnn):
    fcSw = fcSwCollection.find_one({"WWNN": {"$eq": switchwwnn}}, {"MGMT_IP": 1, "SN": 1, "PORTS": 1})
    sn = ""
    doaminIdx2wwpnMap = {}
    doaminIdx2peerwwpnMap = {}
    peerwwpn2wwpnMap = {}
    peerwwpn2domainIdxMap = {}
    peerwwpn2portMap = {}
    if fcSw is not None:
        sn = fcSw.get("SN")
        for port in fcSw.get("PORTS"):
            doaminIdx2wwpnMap[port["DOMAIN_IDX"]] = port["WWPN"]
            peerWwpn = port["PEER_WWPN"]
            if peerWwpn is not None:
                doaminIdx2peerwwpnMap[port["DOMAIN_IDX"]] = peerWwpn
                peerwwpn2wwpnMap[peerWwpn] = port["WWPN"]
                peerwwpn2domainIdxMap[peerWwpn] = port["DOMAIN_IDX"]
                peerwwpn2portMap[peerWwpn] = port["PORT"]

    return (
        sn,
        doaminIdx2wwpnMap,
        doaminIdx2peerwwpnMap,
        peerwwpn2wwpnMap,
        peerwwpn2domainIdxMap,
        peerwwpn2portMap,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--type", default="OS", help="关系类别[OS|NET]")
    args = parser.parse_args()

    (dbclient, db) = AutoExecUtils.getDB()

    try:
        if args.type == "OS" or args.type is None:
            genInsRel(db)
            genDBInsRel(db)
            genOSRel(db)
            genHostRel(db)
            genSwitchRel(db)
            genStorageRel(db)
            genLBRel(db)
        elif args.type == "NET" or args.type is None:
            genSwitchRel(db)
            handleFCSwitchCacadePort(db)
    except Exception as ex:
        print("ERROR: Unknow Error, {}".format(traceback.format_exc()))
        exit(-1)
    finally:
        if dbclient is not None:
            dbclient.close()
